from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableMap, RunnableLambda
from langchain_openai import ChatOpenAI
from config.load_key import load_key

# 提示词模板
prompt_template = ChatPromptTemplate.from_messages([
    ("system", "Translate the following from English into {language}"),
    ("user", "{text}"),
])
# 构建大模型客户端
llm = ChatOpenAI(
    model="Qwen/Qwen2.5-7B-Instruct",
    api_key=load_key("siliconflow_api_key"),
    base_url=load_key("siliconflow_base_url"),
)
# 结果解析器 StrOutputParser 会 AIMessage 转为 str，实际上就是获取 AIMessage.content
output_parser = StrOutputParser()
# 构建链
chain = prompt_template | llm | output_parser
# 调用链
print(chain.invoke({"language": "Chinese", "text": "nice to meet you"}))
"""
chain = prompt_template | llm | output_parser 
print(chain.invoke({"language": "Chinese", "text": "nice to meet you"}))
相当于：
prompt = prompt_template.invoke({"language": "Chinese", "text": "nice to meet you"})
response = llm.invoke(prompt)
print(output_parser.invoke(response))
"""

# 继续构建更复杂的链
analysis_prompt = ChatPromptTemplate.from_template("我应该怎么回答这句话？{talk}。给我一个五个字以内的示例")
analysis_chain = {"talk": chain} | analysis_prompt | llm | output_parser
print(analysis_chain.invoke({"language": "Chinese", "text": "nice to meet you"}))


# 并行执行两个链
# 提示词模板
prompt_template_zh = ChatPromptTemplate.from_messages([
    ("system", "Translate the following from English into Chinese"),
    ("user", "{text}"),
])
prompt_template_fr = ChatPromptTemplate.from_messages([
    ("system", "Translate the following from English into French"),
    ("user", "{text}"),
])
# 构建链
chain_zh = prompt_template_zh | llm | output_parser
chain_fr = prompt_template_fr | llm | output_parser  # 扩展: 翻译为外语时可以使用国外大模型，只需要替换 llm
# 并行执行两个链
parallel_chains = RunnableMap({
    "zh_translation": chain_zh,
    "fr_translation": chain_fr
})
# 合并结果
final_chain = parallel_chains | RunnableLambda(lambda x: f"Chinese: {x['zh_translation']}\n"
                                                         f"French: {x['fr_translation']}")
print(final_chain.invoke({"text": "nice to meet you"}))

# 打印图形化链图
final_chain.get_graph().print_ascii()

"""
注意：只有继承 Runnable 的类 和 只有一个方法的类 才能作为链的节点
链的处理过程是从前往后一个一个处理
"""